from PyQt5.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
                             QLabel, QFileDialog, QSplitter, QSizePolicy, QMessageBox,
                             QDialog, QDialogButtonBox, QListWidget, QListWidgetItem,
                             QSlider, QCheckBox, QColorDialog, QTableWidget, QTableWidgetItem,
                             QComboBox, QTextEdit, QGroupBox, QFormLayout)
from PyQt5.QtCore import Qt, pyqtSlot
from PyQt5.QtGui import QColor, QTextCursor
import numpy as np
import open3d as o3d
from datetime import datetime

from core.point_cloud_viewer import PointCloudViewer
from core.processing_threads import CompletionThread, SegmentationThread, ReconstructionThread
from core.segmentation_methods import SegmentationMethod, PointCloudSegmentor
from core.reconstruction_methods import ReconstructionMethod, PointCloudReconstructor
from ui.success_dialog import SuccessDialog
from ui.waiting_dialog import WaitingDialog
from config import WINDOW_SIZE
import matplotlib.pyplot as plt  # 用于颜色映射
plt.switch_backend('Agg')  # 非交互式后端，避免Qt冲突


# 分割方法选择对话框（保留原有）
class SegmentationMethodDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("选择分割方法")
        self.setFixedSize(300, 200)
        layout = QVBoxLayout()
        self.method_list = QListWidget()
        for method in SegmentationMethod.get_all_methods():
            QListWidgetItem(method, self.method_list)
        self.method_list.setCurrentRow(0)
        layout.addWidget(self.method_list)
        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel, Qt.Horizontal, self)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)
        self.setLayout(layout)
    def get_selected_method(self):
        return self.method_list.currentItem().text()

# 重建方法选择对话框（保留原有）
class ReconstructionMethodDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("选择重建方法")
        self.setFixedSize(300, 200)
        layout = QVBoxLayout()
        self.method_list = QListWidget()
        for method in ReconstructionMethod.get_all_methods():
            QListWidgetItem(method, self.method_list)
        self.method_list.setCurrentRow(0)
        layout.addWidget(self.method_list)
        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel, Qt.Horizontal, self)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)
        self.setLayout(layout)
    def get_selected_method(self):
        return self.method_list.currentItem().text()

# 加载成功对话框（保留原有美化版）
# 加载成功对话框（修复 QColor 无 rgbSwapped() 问题）
class SuccessLoadDialog(QDialog):
    def __init__(self, point_cloud_count, file_ext, single_pc_count, parent=None):
        super().__init__(parent)
        self.setWindowTitle("加载成功")
        self.setFixedSize(400, 220)
        self.setStyleSheet("""
            QDialog { background-color: #f8f9fa; border-radius: 10px; border: 1px solid #e0e0e0; }
            QLabel#title { font-size: 18px; font-weight: bold; color: #2d3748; }
            QLabel#content { font-size: 14px; color: #4a5568; line-height: 1.6; }
            QPushButton { background-color: #4299e1; color: white; border: none; border-radius: 6px; padding: 8px 24px; font-size: 14px; }
            QPushButton:hover { background-color: #3182ce; }
            QPushButton:pressed { background-color: #2b6cb0; }
        """)
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(15)
        top_layout = QHBoxLayout()
        
        # ---------------------- 修复图标绘制逻辑 ----------------------
        icon_label = QLabel()
        # 直接用 16 进制颜色值（蓝色：#007aff），避免 rgbSwapped() 方法
        icon_label.setStyleSheet("background-color: #007aff; border-radius: 32px; width:64px; height:64px;")
        # --------------------------------------------------------------
        
        top_layout.addWidget(icon_label, alignment=Qt.AlignTop)
        text_layout = QVBoxLayout()
        title_label = QLabel("点云加载成功！")
        title_label.setObjectName("title")
        content_label = QLabel(f"共加载 {point_cloud_count} 个点云文件\n格式：.{file_ext}\n当前点云点数：{single_pc_count} 个")
        content_label.setObjectName("content")
        content_label.setWordWrap(True)
        text_layout.addWidget(title_label)
        text_layout.addWidget(content_label)
        top_layout.addLayout(text_layout)
        main_layout.addLayout(top_layout)
        line = QWidget()
        line.setStyleSheet("background-color: #e0e0e0; height:1px;")
        main_layout.addWidget(line)
        btn_layout = QHBoxLayout()
        confirm_btn = QPushButton("确定")
        confirm_btn.clicked.connect(self.accept)
        btn_layout.addWidget(confirm_btn, alignment=Qt.AlignCenter)
        main_layout.addLayout(btn_layout)


# 重建成功对话框（同步修复 QColor 问题）
class SuccessReconstructDialog(QDialog):
    def __init__(self, method_name, vertex_count, triangle_count, parent=None):
        super().__init__(parent)
        self.setWindowTitle("重建成功")
        self.setFixedSize(400, 220)
        self.setStyleSheet("""
            QDialog { background-color: #f8f9fa; border-radius: 10px; border: 1px solid #e0e0e0; }
            QLabel#title { font-size: 18px; font-weight: bold; color: #2d3748; }
            QLabel#content { font-size: 14px; color: #4a5568; line-height: 1.6; }
            QPushButton { background-color: #48bb78; color: white; border: none; border-radius: 6px; padding: 8px 24px; font-size: 14px; }
            QPushButton:hover { background-color: #38a169; }
            QPushButton:pressed { background-color: #2f855a; }
        """)
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(15)
        top_layout = QHBoxLayout()
        
        # ---------------------- 修复图标绘制逻辑 ----------------------
        icon_label = QLabel()
        # 直接用 16 进制颜色值（绿色：#34d399）
        icon_label.setStyleSheet("background-color: #34d399; border-radius: 32px; width:64px; height:64px;")
        # --------------------------------------------------------------
        
        top_layout.addWidget(icon_label, alignment=Qt.AlignTop)
        text_layout = QVBoxLayout()
        title_label = QLabel("点云重建成功！")
        title_label.setObjectName("title")
        content_label = QLabel(f"重建方法：{method_name}\n顶点数：{vertex_count} 个\n三角形数：{triangle_count} 个")
        content_label.setObjectName("content")
        content_label.setWordWrap(True)
        text_layout.addWidget(title_label)
        text_layout.addWidget(content_label)
        top_layout.addLayout(text_layout)
        main_layout.addLayout(top_layout)
        line = QWidget()
        line.setStyleSheet("background-color: #e0e0e0; height:1px;")
        main_layout.addWidget(line)
        btn_layout = QHBoxLayout()
        confirm_btn = QPushButton("确定")
        confirm_btn.clicked.connect(self.accept)
        btn_layout.addWidget(confirm_btn, alignment=Qt.AlignCenter)
        main_layout.addLayout(btn_layout)
# 主窗口类（核心修改）
class PointCloudMainWindow(QMainWindow):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("点云处理系统（增强可视化版）")
        self.setFixedSize(*WINDOW_SIZE)
        self.current_point_clouds = None  # (N,M,3) 点云数据
        self.current_index = 0
        self.max_index = 0
        self.process_thread = None
        self.waiting_dialog = None
        self.current_pcd = None  # 当前 Open3D 点云对象（用于计算几何属性）
        self.current_curvature = None  # 存储曲率数据（用于着色）
        self.current_labels = None  # 存储分割标签（用于着色）
        
        self.init_ui()  # 初始化 UI
         # 关键：初始化按钮状态（未加载点云时，仅加载按钮可用）
        self._set_buttons_enabled(enabled=False)  # 先禁用所有按钮
        self.load_btn.setEnabled(True)  # 单独启用「加载点云」按钮
        self.init_logger()  # 初始化日志
        self.log(f"系统启动成功 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    def init_ui(self):
        """初始化主窗口 UI（新增控制面板、几何属性、日志面板）"""
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)
        main_layout.setContentsMargins(10, 10, 10, 10)
        main_layout.setSpacing(10)

        # ---------------------- 左侧：控制面板（新增）----------------------
        left_panel = QWidget()
        left_panel.setFixedWidth(300)
        left_layout = QVBoxLayout(left_panel)
        left_layout.setSpacing(15)

        # 1. 点云属性调节组
        prop_group = QGroupBox("点云属性调节")
        prop_layout = QFormLayout(prop_group)
        
        # 点大小调节
        self.point_size_slider = QSlider(Qt.Horizontal)
        self.point_size_slider.setRange(1, 10)
        self.point_size_slider.setValue(2)
        self.point_size_slider.setTickPosition(QSlider.TicksBelow)
        prop_layout.addRow(QLabel("点大小："), self.point_size_slider)
        self.point_size_label = QLabel("2")
        prop_layout.addRow(None, self.point_size_label)
        self.point_size_slider.valueChanged.connect(self.on_point_size_changed)

        # 透明度调节
        self.opacity_slider = QSlider(Qt.Horizontal)
        self.opacity_slider.setRange(5, 10)  # 0.5~1.0
        self.opacity_slider.setValue(10)
        self.opacity_slider.setTickPosition(QSlider.TicksBelow)
        prop_layout.addRow(QLabel("透明度："), self.opacity_slider)
        self.opacity_label = QLabel("1.0")
        prop_layout.addRow(None, self.opacity_label)
        self.opacity_slider.valueChanged.connect(self.on_opacity_changed)

        # 背景色选择
        self.bg_color_btn = QPushButton("选择背景色")
        self.bg_color_btn.setStyleSheet("background-color: black; color: white;")
        prop_layout.addRow(QLabel("背景色："), self.bg_color_btn)
        self.bg_color_btn.clicked.connect(self.on_bg_color_changed)

        # 坐标轴显示
        self.axis_checkbox = QCheckBox("显示坐标轴")
        self.axis_checkbox.setChecked(False)
        prop_layout.addRow(self.axis_checkbox)
        self.axis_checkbox.stateChanged.connect(self.on_axis_toggled)

        left_layout.addWidget(prop_group)

        # 2. 颜色模式选择组
        color_group = QGroupBox("颜色模式")
        color_layout = QVBoxLayout(color_group)
        self.color_combo = QComboBox()
        self.color_combo.addItems([
            "原始色（默认蓝）",
            "X坐标着色",
            "Y坐标着色",
            "Z坐标着色",
            "曲率着色（需分割）",
            "分割标签着色（需分割）"
        ])
        color_layout.addWidget(self.color_combo)
        self.color_combo.currentIndexChanged.connect(self.on_color_mode_changed)
        left_layout.addWidget(color_group)

        # 3. 功能按钮（保留原有）
        btn_layout = QVBoxLayout()
        self.load_btn = QPushButton("加载点云")
        self.complete_btn = QPushButton("点云补全")
        self.segment_btn = QPushButton("点云分割")
        self.reconstruct_btn = QPushButton("点云重建")
        self.save_btn = QPushButton("保存结果")
        self.prev_btn = QPushButton("上一个")
        self.next_btn = QPushButton("下一个")
        
        btn_list = [self.load_btn, self.complete_btn, self.segment_btn, self.reconstruct_btn, self.save_btn]
        for btn in btn_list:
            btn.setStyleSheet("padding: 10px; font-size: 14px;")
            btn_layout.addWidget(btn)
        
        # 翻页按钮布局
        nav_layout = QHBoxLayout()
        nav_layout.addWidget(self.prev_btn)
        nav_layout.addWidget(self.next_btn)
        btn_layout.addLayout(nav_layout)
        
        # 绑定按钮事件
        self.load_btn.clicked.connect(self.load_point_cloud)
        self.complete_btn.clicked.connect(self.execute_completion)
        self.segment_btn.clicked.connect(self.execute_segmentation)
        self.reconstruct_btn.clicked.connect(self.execute_reconstruction)
        self.save_btn.clicked.connect(self.save_result)
        self.prev_btn.clicked.connect(self.prev_point_cloud)
        self.next_btn.clicked.connect(self.next_point_cloud)
        
        left_layout.addLayout(btn_layout)
        left_layout.addStretch()

        # ---------------------- 中间：可视化区域（保留原有）----------------------
        mid_panel = QWidget()
        mid_layout = QVBoxLayout(mid_panel)
        mid_layout.setSpacing(10)

        # 原始点云视图
        self.original_title = QLabel("原始点云")
        self.original_title.setAlignment(Qt.AlignCenter)
        self.original_title.setStyleSheet("font-size: 16px; font-weight: bold;")
        self.original_viewer = PointCloudViewer()
        mid_layout.addWidget(self.original_title)
        mid_layout.addWidget(self.original_viewer)

        # 处理结果视图
        self.result_title_label = QLabel("处理结果")
        self.result_title_label.setAlignment(Qt.AlignCenter)
        self.result_title_label.setStyleSheet("font-size: 16px; font-weight: bold;")
        self.result_viewer = PointCloudViewer()
        mid_layout.addWidget(self.result_title_label)
        mid_layout.addWidget(self.result_viewer)

        # ---------------------- 右侧：信息面板（新增）----------------------
        right_panel = QWidget()
        right_panel.setFixedWidth(350)
        right_layout = QVBoxLayout(right_panel)
        right_layout.setSpacing(15)

        # 1. 几何属性表格
        geo_group = QGroupBox("点云几何属性")
        geo_layout = QVBoxLayout(geo_group)
        self.geo_table = QTableWidget(6, 2)  # 6行2列
        self.geo_table.setHorizontalHeaderLabels(["属性", "值"])
        self.geo_table.horizontalHeader().setStretchLastSection(True)
        # 初始化属性行
        self.geo_props = [
            "点云点数",
            "质心(XYZ)",
            "AABB包围盒（最小）",
            "AABB包围盒（最大）",
            "点云密度（点/m³）",
            "表面积（m²）"
        ]
        for i, prop in enumerate(self.geo_props):
            self.geo_table.setItem(i, 0, QTableWidgetItem(prop))
            self.geo_table.setItem(i, 1, QTableWidgetItem("-"))
        geo_layout.addWidget(self.geo_table)
        right_layout.addWidget(geo_group)

        # 2. 日志输出文本框
        log_group = QGroupBox("处理日志")
        log_layout = QVBoxLayout(log_group)
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setStyleSheet("background-color: #f0f0f0; font-family: monospace; font-size: 12px;")
        log_layout.addWidget(self.log_text)
        right_layout.addWidget(log_group)

        # ---------------------- 组装主布局 ----------------------
        main_layout.addWidget(left_panel)
        main_layout.addWidget(mid_panel, stretch=1)
        main_layout.addWidget(right_panel)

        # 禁用初始按钮
        self._set_buttons_enabled(False)

    def init_logger(self):
        """初始化日志系统"""
        self.log_text.clear()
        self.log(f"日志初始化完成 - 支持记录加载、处理、保存等操作")

    def log(self, message):
        """添加日志（带时间戳）"""
        timestamp = datetime.now().strftime("%H:%M:%S")
        self.log_text.append(f"[{timestamp}] {message}")
        # 自动滚动到最新日志
        self.log_text.moveCursor(QTextCursor.End)

    def load_point_cloud(self):
        """加载点云（保留原有逻辑，新增日志和几何属性计算）"""
        file_filter = "点云文件 (*.pt *.pcd *.npy);;PyTorch张量 (*.pt);;PCD格式 (*.pcd);;NumPy数组 (*.npy)"
        file_path, _ = QFileDialog.getOpenFileName(self, "选择点云文件", "", file_filter)
        if not file_path:
            return

        try:
            file_ext = file_path.lower().split('.')[-1]
            self.current_point_clouds = None
            self.max_index = 0

            # 处理不同格式
            if file_ext == 'pt':
                import torch
                data = torch.load(file_path, map_location='cpu')
                if data.dim() == 2:
                    if data.shape[1] == 3:
                        self.current_point_clouds = data.numpy().reshape(1, -1, 3)
                    elif data.shape[0] == 3:
                        self.current_point_clouds = data.permute(1, 0).numpy().reshape(1, -1, 3)
                    else:
                        raise ValueError(f".pt张量形状错误：{data.shape}")
                elif data.dim() == 3:
                    if data.shape[2] == 3:
                        self.current_point_clouds = data.numpy()
                    elif data.shape[1] == 3:
                        self.current_point_clouds = data.transpose(0, 2, 1).numpy()
                    else:
                        raise ValueError(f".pt张量形状错误：{data.shape}")
                else:
                    raise ValueError(f".pt张量维度错误：{data.dim()}")

            elif file_ext == 'pcd':
                pcd = o3d.io.read_point_cloud(file_path)
                if not pcd.has_points():
                    raise ValueError("PCD文件无点云数据")
                point_cloud = np.asarray(pcd.points)
                if point_cloud.shape[1] != 3:
                    raise ValueError(f"PCD点云维度错误：{point_cloud.shape[1]}")
                self.current_point_clouds = point_cloud.reshape(1, -1, 3)
                self.current_pcd = pcd  # 保存Open3D对象

            elif file_ext == 'npy':
                data = np.load(file_path, allow_pickle=False)
                if data.ndim == 2:
                    if data.shape[1] == 3:
                        self.current_point_clouds = data.reshape(1, -1, 3)
                    elif data.shape[0] == 3:
                        self.current_point_clouds = data.T.reshape(1, -1, 3)
                    else:
                        raise ValueError(f".npy数组形状错误：{data.shape}")
                elif data.ndim == 3:
                    if data.shape[2] == 3:
                        self.current_point_clouds = data
                    elif data.shape[1] == 3:
                        self.current_point_clouds = data.transpose(0, 2, 1)
                    else:
                        raise ValueError(f".npy数组形状错误：{data.shape}")
                else:
                    raise ValueError(f".npy数组维度错误：{data.ndim}")

            else:
                raise ValueError(f"不支持的格式：.{file_ext}")

            # 验证格式
            if self.current_point_clouds.ndim != 3 or self.current_point_clouds.shape[2] != 3:
                raise RuntimeError(f"点云格式错误：{self.current_point_clouds.shape}")

            # 更新索引和显示
            self.max_index = self.current_point_clouds.shape[0] - 1
            self.current_index = 0
            self.display_current_point_cloud()  # 显示原始点云
            self.calculate_geometry_props()     # 计算几何属性
            self.update_index_buttons()
            self._set_buttons_enabled(True)

            # 日志和提示
            single_pc_count = self.current_point_clouds[self.current_index].shape[0]
            self.log(f"成功加载点云文件：{file_path}")
            self.log(f"格式：.{file_ext}，总数量：{self.current_point_clouds.shape[0]}，当前点数：{single_pc_count}")
            
            dialog = SuccessLoadDialog(
                point_cloud_count=self.current_point_clouds.shape[0],
                file_ext=file_ext,
                single_pc_count=single_pc_count,
                parent=self
            )
            dialog.exec_()
                    # 加载成功后，启用所有功能按钮


        except Exception as e:
            error_msg = f"加载失败：{str(e)}"
            QMessageBox.critical(self, "加载失败", error_msg)
            self.log(f"加载失败：{str(e)}")

    def display_current_point_cloud(self):
        """显示当前点云（适配颜色模式和属性）"""
        current_pc = self.current_point_clouds[self.current_index]
        # 保存Open3D点云对象（用于后续计算）
        if self.current_pcd is None:
            self.current_pcd = o3d.geometry.PointCloud()
            self.current_pcd.points = o3d.utility.Vector3dVector(current_pc)
        
        # 根据当前颜色模式设置颜色
        color_mode = self.color_combo.currentIndex()
        colors = self.get_color_by_mode(current_pc, color_mode)
        
        # 显示原始点云（应用点大小和透明度）
        self.original_viewer.set_point_cloud(
            current_pc,
            point_colors=colors,
            point_size=self.point_size_slider.value(),
            opacity=self.opacity_slider.value() / 10.0,
            show_axis=self.axis_checkbox.isChecked()
        )
        self.original_title.setText(f"原始点云（索引：{self.current_index+1}/{self.max_index+1}）")

    def get_color_by_mode(self, point_cloud, mode):
        """根据颜色模式生成颜色数组（修复版）"""
        n_points = point_cloud.shape[0]
        colors = None

        if mode == 0:  # 原始色（默认蓝）
            colors = np.tile([0.4, 0.7, 1.0], (n_points, 1))  # 天蓝色

        elif mode == 1:  # X坐标着色
            x = point_cloud[:, 0]
            x_norm = (x - x.min()) / (x.max() - x.min() + 1e-6)  # 归一化到0-1
            colors = plt.cm.viridis(x_norm)[:, :3]  # viridis配色（蓝绿黄渐变）

        elif mode == 2:  # Y坐标着色
            y = point_cloud[:, 1]
            y_norm = (y - y.min()) / (y.max() - y.min() + 1e-6)
            colors = plt.cm.plasma(y_norm)[:, :3]  # plasma配色（紫粉黄渐变）

        elif mode == 3:  # Z坐标着色
            z = point_cloud[:, 2]
            z_norm = (z - z.min()) / (z.max() - z.min() + 1e-6)
            colors = plt.cm.coolwarm(z_norm)[:, :3]  # coolwarm配色（蓝红渐变）

        elif mode == 4:  # 曲率着色（需分割）
            if self.current_curvature is None:
                self.calculate_curvature()  # 自动计算曲率
            curv_norm = (self.current_curvature - self.current_curvature.min()) / (self.current_curvature.max() - self.current_curvature.min() + 1e-6)
            colors = plt.cm.jet(curv_norm)[:, :3]  # jet配色（蓝绿红黄渐变）

        elif mode == 5:  # 分割标签着色（需分割）
            if self.current_labels is None:
                QMessageBox.warning(self, "警告", "请先执行点云分割！")
                colors = np.tile([0.4, 0.7, 1.0], (n_points, 1))  # fallback到原始色
            else:
                max_label = self.current_labels.max()
                color_map = plt.cm.tab10(np.linspace(0, 1, max_label + 1))[:, :3]  # 10种颜色循环
                colors = color_map[self.current_labels]
                colors[self.current_labels == -1] = [0.5, 0.5, 0.5]  # 噪声点灰色

        # 确保颜色数组格式正确（float32，0-1范围）
        if colors is not None:
            colors = colors.astype(np.float32)
        return colors

    def calculate_curvature(self):
        """计算点云曲率（用于着色）"""
        if self.current_pcd is None:
            return
        # 估计法线
        self.current_pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
        points = np.asarray(self.current_pcd.points)
        normals = np.asarray(self.current_pcd.normals)
        
        # 计算曲率
        pcd_tree = o3d.geometry.KDTreeFlann(self.current_pcd)
        curvatures = []
        for i in range(len(points)):
            [k, idx, _] = pcd_tree.search_knn_vector_3d(self.current_pcd.points[i], 30)
            neighbor_normals = normals[idx]
            avg_normal = np.mean(neighbor_normals, axis=0)
            curvature = 1 - np.abs(np.dot(normals[i], avg_normal))
            curvatures.append(curvature)
        
        self.current_curvature = np.array(curvatures)
        self.log(f"成功计算点云曲率，曲率范围：[{self.current_curvature.min():.4f}, {self.current_curvature.max():.4f}]")

    def calculate_geometry_props(self):
        """计算点云几何属性并更新UI"""
        if self.current_pcd is None or self.current_point_clouds is None:
            return
        
        current_pc = self.current_point_clouds[self.current_index]
        n_points = current_pc.shape[0]
        
        # 1. 点云点数
        self.geo_table.setItem(0, 1, QTableWidgetItem(f"{n_points}"))
        
        # 2. 质心
        centroid = np.mean(current_pc, axis=0)
        self.geo_table.setItem(1, 1, QTableWidgetItem(f"{centroid[0]:.6f}, {centroid[1]:.6f}, {centroid[2]:.6f}"))
        
        # 3. AABB包围盒
        aabb_min = np.min(current_pc, axis=0)
        aabb_max = np.max(current_pc, axis=0)
        self.geo_table.setItem(2, 1, QTableWidgetItem(f"{aabb_min[0]:.6f}, {aabb_min[1]:.6f}, {aabb_min[2]:.6f}"))
        self.geo_table.setItem(3, 1, QTableWidgetItem(f"{aabb_max[0]:.6f}, {aabb_max[1]:.6f}, {aabb_max[2]:.6f}"))
        
        # 4. 点云密度（简化：点数 / 包围盒体积）
        aabb_volume = (aabb_max[0] - aabb_min[0]) * (aabb_max[1] - aabb_min[1]) * (aabb_max[2] - aabb_min[2])
        density = n_points / (aabb_volume + 1e-6)  # 避免除零
        self.geo_table.setItem(4, 1, QTableWidgetItem(f"{density:.2f}"))
        
        # 5. 表面积（需重建后计算）
        self.geo_table.setItem(5, 1, QTableWidgetItem("-"))  # 重建后更新
        
        self.log(f"成功计算几何属性：点数={n_points}，质心={centroid.round(6)}，密度={density:.2f}点/m³")

    # ---------------------- 新增：属性调节槽函数 ----------------------
    @pyqtSlot(int)
    def on_point_size_changed(self, value):
        """点大小调节"""
        self.point_size_label.setText(str(value))
        if self.current_point_clouds is not None:
            self.original_viewer.update_point_size(value)
            if self.result_viewer.point_cloud is not None:
                self.result_viewer.update_point_size(value)

    @pyqtSlot(int)
    def on_opacity_changed(self, value):
        """透明度调节（0.5~1.0）"""
        opacity = value / 10.0
        self.opacity_label.setText(f"{opacity:.1f}")
        if self.current_point_clouds is not None:
            self.original_viewer.update_opacity(opacity)
            if self.result_viewer.point_cloud is not None:
                self.result_viewer.update_opacity(opacity)

    @pyqtSlot()
    def on_bg_color_changed(self):
        """背景色选择"""
        color = QColorDialog.getColor(Qt.black, self, "选择背景色")
        if color.isValid():
            self.bg_color_btn.setStyleSheet(f"background-color: {color.name()}; color: white;")
            r = color.redF()
            g = color.greenF()
            b = color.blueF()
            self.original_viewer.update_bg_color(r, g, b)
            self.result_viewer.update_bg_color(r, g, b)
            self.log(f"背景色更新为：RGB({r:.2f}, {g:.2f}, {b:.2f})")

    @pyqtSlot(int)
    def on_axis_toggled(self, state):
        """坐标轴显示切换"""
        show_axis = state == Qt.Checked
        self.original_viewer.update_show_axis(show_axis)
        self.result_viewer.update_show_axis(show_axis)
        self.log(f"坐标轴显示：{'开启' if show_axis else '关闭'}")

    @pyqtSlot(int)
    def on_color_mode_changed(self, index):
        """颜色模式切换（修复：同步结果视图）"""
        mode_name = self.color_combo.itemText(index)
        self.log(f"颜色模式切换为：{mode_name}")
        if self.current_point_clouds is not None:
            # 刷新原始点云颜色
            self.display_current_point_cloud()
            
            # 刷新结果点云颜色（修复核心）
            if self.result_viewer.point_cloud is not None:
                result_pc = self.result_viewer.point_cloud
                colors = self.get_color_by_mode(result_pc, index)
                # 重新设置结果点云（带新颜色）
                self.result_viewer.set_point_cloud(
                    result_pc, 
                    point_colors=colors,
                    point_size=self.point_size_slider.value(),
                    opacity=self.opacity_slider.value() / 10.0
                )
            # 刷新网格颜色（如果存在）
            elif self.result_viewer.mesh is not None:
                # 网格颜色暂时保持不变，如需支持可扩展
                pass

    # ---------------------- 原有功能保持不变（补全、分割、重建）----------------------
    def execute_completion(self):
        if self.current_point_clouds is None:
            QMessageBox.warning(self, "警告", "请先加载点云！")
            return
        self._set_buttons_enabled(False)
        current_pc = self.current_point_clouds[self.current_index]

        # 创建等待对话框，并绑定取消信号
        self.waiting_dialog = WaitingDialog(message="正在执行点云补全...", parent=self)
        # 关键：点击取消按钮时，调用 cancel_completion 终止线程
        self.waiting_dialog.cancel_clicked.connect(self.cancel_completion)
        self.waiting_dialog.show()

        try:
            # 创建补全线程
            self.process_thread = CompletionThread(current_pc)
            self.process_thread.completion_finished.connect(self._handle_task_success)
            self.process_thread.completion_error.connect(self._handle_task_error)
            self.process_thread.start()
            self.log("开始执行点云补全...")
        except Exception as e:
            error_msg = f"线程创建失败：{str(e)}"
            QMessageBox.critical(self, "错误", error_msg)
            self.log(error_msg)
            self.waiting_dialog.close()
            self._set_buttons_enabled(True)

    def cancel_completion(self):
        """终止点云补全线程"""
        if self.process_thread and self.process_thread.isRunning():
            # 终止线程（强制结束）
            self.process_thread.terminate()
            self.process_thread.wait()  # 等待线程资源释放
            self.log("点云补全已取消")
            QMessageBox.information(self, "已取消", "点云补全处理已终止")
        # 恢复界面状态
        self._set_buttons_enabled(True)
        # 确保等待窗口关闭
        if self.waiting_dialog:
            self.waiting_dialog.close()

    def execute_segmentation(self):
        if self.current_point_clouds is None:
            QMessageBox.warning(self, "警告", "请先加载点云！")
            return
        method_dialog = SegmentationMethodDialog(self)
        if method_dialog.exec_() != QDialog.Accepted:
            return
        selected_method = method_dialog.get_selected_method()
        self._set_buttons_enabled(False)
        current_pc = self.current_point_clouds[self.current_index]
        self.waiting_dialog = WaitingDialog(message=f"正在执行{selected_method}...", parent=self)
        self.waiting_dialog.show()
        self.process_thread = SegmentationThread(current_pc, selected_method)
        self.process_thread.segmentation_finished.connect(self._handle_segmentation_success)
        self.process_thread.segmentation_error.connect(self._handle_task_error)
        self.process_thread.start()
        self.log(f"开始执行{selected_method}...")

    def execute_reconstruction(self):
        if self.current_point_clouds is None:
            QMessageBox.warning(self, "警告", "请先加载点云！")
            return
        
        # ---------------------- 新增：弹出重建方法选择对话框 ----------------------
        method_dialog = ReconstructionMethodDialog(self)  # 使用已定义的对话框
        if method_dialog.exec_() != QDialog.Accepted:  # 用户点击取消
            return
        selected_method = method_dialog.get_selected_method()  # 获取选择的方法
        # -------------------------------------------------------------------------
        
        self._set_buttons_enabled(False)
        current_pc = self.current_point_clouds[self.current_index]
        # 等待对话框显示具体方法名称
        self.waiting_dialog = WaitingDialog(message=f"正在执行{selected_method}...", parent=self)
        self.waiting_dialog.cancel_clicked.connect(self.cancel_reconstruction)
        self.waiting_dialog.show()

        try:
            # 修复：传入 point_cloud_data 和 reconstruct_method 两个参数
            self.process_thread = ReconstructionThread(
                point_cloud_data=current_pc,
                reconstruct_method=selected_method  # 补充从对话框获取的方法
            )
            # 注意：重建线程的成功信号需要绑定到专门的处理函数（而非通用的 _handle_task_success）
            self.process_thread.reconstruction_finished.connect(self._handle_reconstruction_success)
            self.process_thread.reconstruction_error.connect(self._handle_task_error)
            self.process_thread.start()
            self.log(f"开始执行{selected_method}...")  # 日志显示具体方法
        except Exception as e:
            self._set_buttons_enabled(True)
            self.waiting_dialog.close()
            QMessageBox.critical(self, "错误", f"线程创建失败：{str(e)}")

    def _handle_segmentation_success(self, points, colors, labels):
        if self.waiting_dialog:
            self.waiting_dialog.close()
        self.current_labels = labels  # 保存标签用于着色
        self.result_viewer.set_point_cloud(points, colors)
        max_label = labels.max()
        part_count = max_label + 1 if max_label >= 0 else 0
        noise_count = np.sum(labels == -1)
        self.result_title_label.setText(f"分割结果（{self.process_thread.segment_method}）- 部件数：{part_count}，噪声点：{noise_count}")
        self.log(f"分割成功：方法={self.process_thread.segment_method}，部件数={part_count}，噪声点={noise_count}")
        
        dialog = SuccessDialog(title="分割成功", content=f"成功使用{self.process_thread.segment_method}分割出 {part_count} 个部件\n噪声点数量：{noise_count}", parent=self)
        dialog.exec_()
        self._set_buttons_enabled(True)

    def _handle_reconstruction_success(self, mesh, vertices, method_name):
        if self.waiting_dialog:
            self.waiting_dialog.close()
        self.result_viewer.set_mesh(mesh)
        vertex_count = len(vertices)
        triangle_count = len(mesh.triangles)
        self.result_title_label.setText(f"重建结果（{method_name}）- 顶点数：{vertex_count}，三角形数：{triangle_count}")
        
        # 更新几何属性中的表面积
        surface_area = mesh.get_surface_area()
        self.geo_table.setItem(5, 1, QTableWidgetItem(f"{surface_area:.6f}"))
        self.log(f"重建成功：方法={method_name}，顶点数={vertex_count}，三角形数={triangle_count}，表面积={surface_area:.6f}m²")
        
        dialog = SuccessReconstructDialog(method_name=method_name, vertex_count=vertex_count, triangle_count=triangle_count, parent=self)
        dialog.exec_()
        self._set_buttons_enabled(True)

    def _handle_task_success(self, result, title):
        """任务成功后的处理（通用）"""
        if self.waiting_dialog:
            self.waiting_dialog.close()
        # 显示结果（区分点云和网格）
        if isinstance(result, np.ndarray):
            # 点云结果（补全/分割）
            self.result_viewer.set_point_cloud(result)
        else:
            # 网格结果（重建）
            self.result_viewer.set_mesh(result)
        self.result_title_label.setText(title)
        self.log(f"{title}成功")
        SuccessDialog(title=title, content=f"{title}成功！", parent=self).exec_()
        self._set_buttons_enabled(True)  # 关键：恢复按钮可用

    def _handle_task_error(self, error):
        """任务失败后的处理（通用）"""
        if self.waiting_dialog:
            self.waiting_dialog.close()
        QMessageBox.critical(self, "处理失败", error)
        self.log(f"处理失败：{error}")
        self._set_buttons_enabled(True)  # 关键：恢复按钮可用

    def save_result(self):
        if self.result_viewer.point_cloud is None and self.result_viewer.mesh is None:
            QMessageBox.warning(self, "警告", "无处理结果可保存！")
            return

        # 选择保存路径（支持点云和网格）
        file_path, _ = QFileDialog.getSaveFileName(
            self, 
            "保存结果", 
            "", 
            "PCD文件 (*.pcd);;PLY文件 (*.ply);;NumPy数组 (*.npy)"  # 增加PLY格式支持网格
        )
        if not file_path:
            return

        # 确保文件带后缀
        if not file_path.endswith(('.pcd', '.ply', '.npy')):
            file_path += '.pcd'  # 默认保存为PCD

        try:
            if file_path.endswith('.pcd') or file_path.endswith('.ply'):
                # 处理点云保存
                if self.result_viewer.point_cloud is not None:
                    pcd = o3d.geometry.PointCloud()
                    pcd.points = o3d.utility.Vector3dVector(self.result_viewer.point_cloud)
                    if self.result_viewer.point_colors is not None:
                        pcd.colors = o3d.utility.Vector3dVector(self.result_viewer.point_colors)
                    # 修复：删除 overwrite=True，Open3D 默认覆盖
                    o3d.io.write_point_cloud(file_path, pcd)  # PCD/PLY通用
                # 处理网格保存（重建结果）
                elif self.result_viewer.mesh is not None:
                    # 网格保存用 write_triangle_mesh
                    o3d.io.write_triangle_mesh(file_path, self.result_viewer.mesh)

            elif file_path.endswith('.npy'):
                # 保存为NumPy数组
                if self.result_viewer.mesh is not None:
                    data = np.asarray(self.result_viewer.mesh.vertices)
                else:
                    data = self.result_viewer.point_cloud
                np.save(file_path, data)

            self.log(f"成功保存结果到：{file_path}")
            QMessageBox.information(self, "保存成功", f"结果已保存到：\n{file_path}")

        except Exception as e:
            error_msg = f"保存失败：{str(e)}"
            QMessageBox.critical(self, "保存失败", error_msg)
            self.log(error_msg)

    def prev_point_cloud(self):
        if self.current_index > 0:
            self.current_index -= 1
            self.display_current_point_cloud()
            self.calculate_geometry_props()
            self.log(f"切换到上一个点云，索引：{self.current_index}")

    def next_point_cloud(self):
        if self.current_index < self.max_index:
            self.current_index += 1
            self.display_current_point_cloud()
            self.calculate_geometry_props()
            self.log(f"切换到下一个点云，索引：{self.current_index}")

    def update_index_buttons(self):
        self.prev_btn.setEnabled(self.current_index > 0)
        self.next_btn.setEnabled(self.current_index < self.max_index)

    def _set_buttons_enabled(self, enabled):
        """统一控制所有功能按钮的状态"""
        self.load_btn.setEnabled(enabled)        # 加载点云按钮
        self.segment_btn.setEnabled(enabled)     # 分割按钮
        self.complete_btn.setEnabled(enabled)    # 补全按钮
        self.reconstruct_btn.setEnabled(enabled) # 重建按钮
        self.save_btn.setEnabled(enabled)        # 保存按钮
        # 其他按钮（如颜色模式、透明度等）也应在此控制
        self.color_combo.setEnabled(enabled)
        self.opacity_slider.setEnabled(enabled)

    def cancel_reconstruction(self):
        """取消点云重建"""
        if self.process_thread and self.process_thread.isRunning():
            self.process_thread.terminate()
            self.process_thread.wait()
            self.log("点云重建已取消")
        self._set_buttons_enabled(True)  # 恢复按钮
        if self.waiting_dialog:
            self.waiting_dialog.close()